package org.wenshu.ai;

import dev.langchain4j.data.document.Metadata;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.embedding.onnx.allminilml6v2.AllMiniLmL6V2EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingSearchResult;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.filter.Filter;
import dev.langchain4j.store.embedding.milvus.MilvusEmbeddingStore;

import java.util.Map;

import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;

public class MilvusWithMetadataExample {

    public static void main(String[] args) {
        EmbeddingStore<TextSegment> embeddingStore = MilvusEmbeddingStore
            .builder()
            .host("localhost").port(19530)
            .collectionName("test_collection2")
            .dimension(384)
            .build();

        EmbeddingModel embeddingModel = new AllMiniLmL6V2EmbeddingModel();

        TextSegment segment1 = TextSegment.from("I like football.", Metadata.metadata("userId", "1"));
        Embedding embedding1 = embeddingModel.embed(segment1).content();
        embeddingStore.add(embedding1, segment1);

        TextSegment segment2 = TextSegment.from("I like basketball.", Metadata.metadata("userId", "2"));
        Embedding embedding2 = embeddingModel.embed(segment2).content();
        embeddingStore.add(embedding2, segment2);

        TextSegment ddlSegment = TextSegment.from("The following columns are in the Actions table in the def database: | | TABLE_CATALOG | TABLE_SCHEMA | TABLE_NAME | COLUMN_NAME | DATA_TYPE | COLUMN_COMMENT | |---:|:----------------|:---------------|:-------------|:------------------|:------------|:-----------------| | 48 | def | NBA | Actions | GameId | int | | | 49 | def | NBA | Actions | TeamId | int | | | 50 | def | NBA | Actions | PlayerId | int | | | 51 | def | NBA | Actions | Minutes | int | | | 52 | def | NBA | Actions | FieldGoalsMade | int | | | 53 | def | NBA | Actions | FieldGoalAttempts | int | | | 54 | def | NBA | Actions | 3PointsMade | int | | | 55 | def | NBA | Actions | 3PointAttempts | int | | | 56 | def | NBA | Actions | FreeThrowsMade | int | | | 57 | def | NBA | Actions | FreeThrowAttempts | int | | | 58 | def | NBA | Actions | PlusMinus | int | | | 59 | def | NBA | Actions | OffensiveRebounds | int | | | 60 | def | NBA | Actions | DefensiveRebounds | int | | | 61 | def | NBA | Actions | TotalRebounds | int | | | 62 | def | NBA | Actions | Assists | int | | | 63 | def | NBA | Actions | PersonalFouls | int | | | 64 | def | NBA | Actions | Steals | int | | | 65 | def | NBA | Actions | Turnovers | int | | | 66 | def | NBA | Actions | BlockedShots | int | | | 67 | def | NBA | Actions | BlocksAgainst | int | | | 68 | def | NBA | Actions | Points | int | | | 69 | def | NBA | Actions | Starter | int | |",
            Metadata.from(Map.of("db", "NBA","table","Actions")));
        Embedding embedding = embeddingModel.embed(ddlSegment).content();
        embeddingStore.add(embedding, ddlSegment);

      Filter filter = metadataKey("db").isEqualTo("NBA").and(metadataKey("table").isEqualTo("Actions"));
      embeddingStore.removeAll(filter);
      EmbeddingSearchRequest searchRequest = EmbeddingSearchRequest
          .builder().queryEmbedding( embeddingModel.embed("NBA table").content())
          .filter(filter)
          .build();

      embeddingStore.search(searchRequest).matches().forEach(System.out::println);

      Embedding queryEmbedding = embeddingModel.embed("What is your favourite sport?").content();

        // search for user 1

        Filter onlyForUser1 = metadataKey("userId").isEqualTo("1");

        EmbeddingSearchRequest embeddingSearchRequest1 = EmbeddingSearchRequest
            .builder()
            .queryEmbedding(queryEmbedding)
            .filter(onlyForUser1)
            .build();

        EmbeddingSearchResult<TextSegment> embeddingSearchResult1 = embeddingStore.search(embeddingSearchRequest1);
        EmbeddingMatch<TextSegment> embeddingMatch1 = embeddingSearchResult1.matches().get(0);

        System.out.println(embeddingMatch1.score());
        System.out.println(embeddingMatch1.embedded().text());

        // search for user 2

        Filter onlyForUser2 = metadataKey("userId").isEqualTo("2");

        EmbeddingSearchRequest embeddingSearchRequest2 = EmbeddingSearchRequest
            .builder()
            .queryEmbedding(queryEmbedding)
            .filter(onlyForUser2)
            .build();

        EmbeddingSearchResult<TextSegment> embeddingSearchResult2 = embeddingStore.search(embeddingSearchRequest2);
        EmbeddingMatch<TextSegment> embeddingMatch2 = embeddingSearchResult2.matches().get(0);

        System.out.println(embeddingMatch2.score());
        System.out.println(embeddingMatch2.embedded().text());

    }
}
